import keras as keras
import numpy as np
import cv2
from keras.preprocessing.image import img_to_array
import matplotlib.pyplot as plt


def img_tool(file_path):
    img = cv2.imread(file_path)
    img = cv2.resize(img, (28, 28))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = img_to_array(img)
    img = np.expand_dims(img, axis=0)
    return img


(x_train, y_train), (x_valid, y_valid) = keras.datasets.mnist.load_data()
assert x_train.shape == (60000, 28, 28)
assert x_valid.shape == (10000, 28, 28)
assert y_train.shape == (60000,)
assert y_valid.shape == (10000,)

y_train_cate = keras.utils.to_categorical(y_train, 10)  # 将y_train转为one-hot编码
y_valid_cate = keras.utils.to_categorical(y_valid, 10)  # 将y_valid转为one-hot编码

# step1: use sequential
model = keras.models.Sequential()

# step2: add layer
model.add(keras.layers.Flatten(input_shape=(x_train.shape[1], x_train.shape[2])))
model.add(keras.layers.Dense(units=784, activation="relu", input_dim=784))
model.add(keras.layers.Dense(units=10, activation="softmax"))

# step3: compile model
model.compile(optimizer="Adam", loss='categorical_crossentropy', metrics=['accuracy'])  # loss使用categorical_crossentropy

print("model:")
model.summary()

# step4: train
model.fit(x_train, y_train_cate, batch_size=64, epochs=3)

# step5: evaluate model
model.evaluate(x_valid, y_valid_cate)  # 评估时使用categorical_crossentropy

# save model
# model.save('keras_mnist.h5')

# img = x_valid[0]
# img = np.reshape(img, (-1, 28, 28))
# 图片处理
img = cv2.imread("44.png")  # 读取图片
img = cv2.resize(img, (28, 28))  # 重置图片大小
img = img[:, :, 0:1]
plt.imshow(img[:, :, ::-1])

img = np.reshape(img, (-1, 28, 28))

output = model.predict(img)
print(output)
predict_num = np.argmax(output, axis=1)
print("predict num is ", predict_num)
